# coding=utf-8
from textrank4zh import TextRank4Keyword, TextRank4Sentence, util
import jieba.analyse
from snownlp import SnowNLP
import pandas as pd
import numpy as np

# 关键词抽取
def keywords_extraction(text, stop_words_file=None
                        , vertex_source='all_filters', edge_source='no_stop_words', num=6, word_min_len=1):
    '''
    关键词抽取
    :param text: 文档来源
    :param stop_words_file: 停用词词典来源
    :param vertex_source:  选择使用words_no_filter, words_no_stop_words, words_all_filters中的哪一个来构造pagerank对应的图中的节点
                           -- 默认值为`'all_filters'`，可选值为`'no_filter', 'no_stop_words', 'all_filters'
    :param edge_source:    选择使用words_no_filter, words_no_stop_words, words_all_filters中的哪一个来构造pagerank对应的图中的节点之间的边
                           -- 默认值为`'no_stop_words'`，可选值为`'no_filter', 'no_stop_words', 'all_filters'`。边的构造要结合`window`参数
    :param num: 返回关键词数量
    :param word_min_len: 词的最小长度，默认值为1
    :return: 关键词list，每个元素为dict，存储了分词及其权重------ TextRank4Keyword().analyze()

    keywords=[{'word': '小行星', 'weight': 0.05808441467341854}, {'word': '天文台', 'weight': 0.057216537757425114}, ……]
    '''
    tr4w = TextRank4Keyword(allow_speech_tags=['n', 'nr', 'nrfg', 'ns', 'nt', 'nz'], stop_words_file=stop_words_file)
    # allow_speech_tags   --运行的词性列表，用于过滤某些词性的词
    # 这里好像是只取这些个名词
    tr4w.analyze(text=text, window=2, lower=True, vertex_source=vertex_source, edge_source=edge_source,
                 pagerank_config={'alpha': 0.85, })
    '''
    # text    --  文本内容，字符串
    # window  --  窗口大小，int，用来构造单词之间的边。默认值为2
    # lower   --  是否将英文文本转换为小写，默认值为False
    # vertex_source  -- 选择使用words_no_filter, words_no_stop_words, words_all_filters中的哪一个来构造pagerank对应的图中的节点
    #                -- 默认值为`'all_filters'`，可选值为`'no_filter', 'no_stop_words', 'all_filters'
    # edge_source  -- 选择使用words_no_filter, words_no_stop_words, words_all_filters中的哪一个来构造pagerank对应的图中的节点之间的边
    #              -- 默认值为`'no_stop_words'`，可选值为`'no_filter', 'no_stop_words', 'all_filters'`。边的构造要结合`window`参数
    '''


    # pagerank_config  -- pagerank算法参数配置，阻尼系数为0.85
    keywords = tr4w.get_keywords(num=num, word_min_len=word_min_len)
    '''
    # num           --  返回关键词数量
    # word_min_len  --  词的最小长度，默认值为1
    '''

    return keywords


# 关键短语抽取
def keyphrases_extraction(text):
    tr4w = TextRank4Keyword()
    tr4w.analyze(text=text, window=2, lower=True, vertex_source='all_filters', edge_source='no_stop_words',
                 pagerank_config={'alpha': 0.85, })
    keyphrases = tr4w.get_keyphrases(keywords_num=6, min_occur_num=1)
    # keywords_num    --  抽取的关键词数量
    # min_occur_num   --  关键短语在文中的最少出现次数
    return keyphrases


# 关键句抽取
def keysentences_extraction(text,
                            stop_words_file=None, allow_speech_tags=util.allow_speech_tags,
                            source='all_filters',
                            num=3):
    tr4s = TextRank4Sentence(stop_words_file=stop_words_file, allow_speech_tags=allow_speech_tags)
    # stop_words_file     --★停用词词典
    # allow_speech_tags   --运行的词性列表，用于过滤某些词性的词
    tr4s.analyze(text, lower=True, source=source)   # 根据（相同的关键词）关键程度定义各个句子的权值，再进行pageRank算法，然后对句子排序
    # text    -- 文本内容，字符串
    # lower   -- 是否将英文文本转换为小写，默认值为False
    # source  -- ★选择使用words_no_filter, words_no_stop_words, words_all_filters中的哪一个来生成句子之间的相似度。
    # 		  -- 默认值为`'all_filters'`，可选值为`'no_filter', 'no_stop_words', 'all_filters'
    # sim_func -- 指定计算句子相似度的函数  # 两个句子的相似度比较（按照相同的关键词作为关键程度）

    # 获取最重要的num个长度大于等于sentence_min_len的句子用来生成摘要
    keysentences = tr4s.get_key_sentences(num=num, sentence_min_len=6)
    return keysentences


def keywords_textrank(text):
    keywords = jieba.analyse.textrank(text, topK=6)
    return keywords


if __name__ == "__main__":
    text = "来源：中国科学报本报讯（记者肖洁）又有一位中国科学家喜获小行星命名殊荣！4月19日下午，中国科学院国家天文台在京举行“周又元星”颁授仪式，" \
           "我国天文学家、中国科学院院士周又元的弟子与后辈在欢声笑语中济济一堂。国家天文台党委书记、" \
           "副台长赵刚在致辞一开始更是送上白居易的诗句：“令公桃李满天下，何须堂前更种花。”" \
           "据介绍，这颗小行星由国家天文台施密特CCD小行星项目组于1997年9月26日发现于兴隆观测站，" \
           "获得国际永久编号第120730号。2018年9月25日，经国家天文台申报，" \
           "国际天文学联合会小天体联合会小天体命名委员会批准，国际天文学联合会《小行星通报》通知国际社会，" \
           "正式将该小行星命名为“周又元星”。"
    # 关键词抽取
    print('关键词抽取:')
    keywords = keywords_extraction(text)
    print(keywords)

    # 关键短语抽取
    print('关键短句抽取:')
    keyphrases = keyphrases_extraction(text)
    print(keyphrases)

    # 关键句抽取
    print('关键句抽取:')
    keysentences = keysentences_extraction(text)
    print(keysentences)